# version:1.1.2312.9221-mp
import gxipy as gx
import time
from ctypes import addressof, c_ubyte
from gxipy.gxidef import GxPixelFormatEntry, DxValidBit
from gxipy.ImageFormatConvert import ImageFormatConvert
import numpy as np
import cv2
import multiprocessing
import os

def get_best_valid_bits(pixel_format):
    vb = DxValidBit.BIT0_7
    # （同原代码，其它格式映射略）
    return vb

def convert_to_bgr(raw_image, img_conv):
    """Convert RawImage → RGB8 numpy → BGR"""
    img_conv.set_dest_format(GxPixelFormatEntry.RGB8)
    img_conv.set_valid_bits(get_best_valid_bits(raw_image.get_pixel_format()))
    buf_size = img_conv.get_buffer_size_for_conversion(raw_image)
    buf = (c_ubyte * buf_size)()
    ptr = addressof(buf)
    img_conv.convert(raw_image, ptr, buf_size, False)
    arr = np.frombuffer(buf, dtype=np.uint8, count=buf_size)
    h, w = raw_image.get_height(), raw_image.get_width()
    rgb = arr.reshape(h, w, 3)
    return cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR)

def camera_worker(cam_idx, trigger_evt, stop_evt, out_dir, log_path):
    """
    子进程：打开 cam_idx 相机，创建 ImageFormatConvert，
    循环等待 trigger_evt，然后软件触发→抓一帧→保存 PNG + 记录 TXT。
    """
    # 准备输出目录与日志
    os.makedirs(out_dir, exist_ok=True)
    log = open(log_path, "w", encoding="utf-8")
    log.write("frame_idx\ttimestamp\n")
    # 打开相机
    dm = gx.DeviceManager()
    num, info = dm.update_all_device_list()
    if num < cam_idx:
        print(f"Camera {cam_idx} not found"); return
    cam = dm.open_device_by_index(cam_idx)
    feat = cam.get_remote_device_feature_control()
    # 软触发 & 曝光
    feat.get_enum_feature("TriggerMode").set("On")
    feat.get_enum_feature("TriggerSource").set("Software")
    feat.get_float_feature("ExposureTime").set(70000.0)
    # 转换器
    img_conv = dm.create_image_format_convert()
    # 数据流 & 触发命令
    ds = cam.data_stream[0]
    trig = feat.get_command_feature("TriggerSoftware")
    cam.stream_on()

    idx = 0
    try:
        while not stop_evt.is_set():
            trigger_evt.wait()  # 等待主进程触发
            if stop_evt.is_set(): break
            trig.send_command()
            raw = ds.get_image()
            if raw:
                frame = convert_to_bgr(raw, img_conv)
                # 存 PNG
                fn = os.path.join(out_dir, f"frame_{idx:06d}.png")
                cv2.imwrite(fn, frame)
                # 记日志
                ts = time.time()
                log.write(f"{idx}\t{ts:.6f}\n")
                idx += 1
            trigger_evt.clear()
    finally:
        log.close()
        cam.stream_off()
        cam.close_device()
        print(f"Camera {cam_idx} worker exit.")

def main():
    mgr = multiprocessing.Manager()
    trigger_evt = mgr.Event()
    stop_evt = mgr.Event()

    # 两个子进程分别对应设备索引 1、2
    p1 = multiprocessing.Process(
        target=camera_worker,
        args=(1, trigger_evt, stop_evt, "frames_cam1", "timestamps_cam1.txt")
    )
    p2 = multiprocessing.Process(
        target=camera_worker,
        args=(2, trigger_evt, stop_evt, "frames_cam2", "timestamps_cam2.txt")
    )
    p1.start(); p2.start()

    print("按 ‘q’ 退出，同步采集开始…")
    try:
        while True:
            trigger_evt.set()                # 同步触发两台相机
            time.sleep(0.005)                 # 控制触发频率 ~30 FPS
            if cv2.waitKey(1) & 0xFF == ord('q'):
                break
    except KeyboardInterrupt:
        pass
    finally:
        stop_evt.set()                      # 通知子进程退出
        trigger_evt.set()                   # 解锁等待
        p1.join(); p2.join()
        cv2.destroyAllWindows()
        print("采集结束。")

if __name__ == "__main__":
    main()
